"""
LSTM_MacroCycle Strategy Implementation
Uses LSTM model to predict macroeconomic cycles and adjust asset allocation accordingly
"""

from zipline.api import order_target_percent, symbol, schedule_function, get_datetime
from zipline.api import date_rules, time_rules
import numpy as np
import pandas as pd
import numpy as np
import pandas as pd

def initialize(context):
    """Initialize strategy parameters and schedules"""
    # Set LSTM model parameters
    context.lookback = 60
    context.confidence_threshold = 0.7
    context.current_cycle = "RECOVERY"  # Default starting cycle
    
    # Initialize mock LSTM model (replace with actual trained model)
    context.model = MockLSTMModel()
    
    # Economic data tracking
    context.economic_data = pd.DataFrame(columns=['pmi', 'cpi'])
    
    # Set scheduler for weekly predictions (Friday market close)
    # Weekly predictions (Friday market close)
    schedule_function(
        predict_macro_cycle,
        date_rule='week_end',
        time_rule='market_close'
    )
    
    # Monthly rebalancing (first day of month)
    schedule_function(
        rebalance,
        date_rule='month_start',
        time_rule='market_open'
    )
    
    # Define asset symbols
    context.stocks = symbol('SPY')  # S&P 500 ETF
    context.bonds = symbol('TLT')   # 20+ Year Treasury Bond ETF
    context.gold = symbol('GLD')    # Gold ETF
    context.cash = symbol('SHY')    # Short-term Treasury ETF (cash equivalent)

def handle_data(context, data):
    """Collect daily economic data"""
    # In a real implementation, this would fetch actual economic indicators
    # For simulation purposes, we'll generate mock data
    today = get_datetime().normalize()
    
    # Generate random economic data (replace with real data source)
    new_data = pd.DataFrame({
        'date': [today],
        'pmi': [np.random.uniform(45, 60)],  # PMI typically 45-60
        'cpi': [np.random.uniform(1.5, 4.0)] # CPI typically 1.5-4.0%
    }).set_index('date')
    
    # Store economic data
    context.economic_data = pd.concat([context.economic_data, new_data])

def predict_macro_cycle(context, data):
    """Predict macroeconomic cycle using LSTM model"""
    if len(context.economic_data) < context.lookback:
        return  # Not enough data yet
    
    # Prepare input sequence (last N days of economic data)
    sequence = context.economic_data[-context.lookback:][['pmi', 'cpi']].values
    
    # Get prediction from model
    prediction, confidence = context.model.predict(sequence)
    
    # Update cycle state if confidence is high enough
    if confidence > context.confidence_threshold:
        context.current_cycle = decode_prediction(prediction)
        print(f"Cycle updated to {context.current_cycle} with confidence {confidence:.2f}")

def rebalance(context, data):
    """Rebalance portfolio based on current economic cycle"""
    # Risk control: Maintain at least 10% cash
    cash_weight = 0.1
    
    # Allocate assets based on current economic cycle
    if context.current_cycle == "RECOVERY":
        # 70% stocks, 20% bonds, 10% cash
        allocate_assets(context, data, [0.7, 0.2, 0.0], cash_weight)
    elif context.current_cycle == "OVERHEAT":
        # 50% commodities, 30% stocks, 10% cash (commodities not implemented)
        allocate_assets(context, data, [0.3, 0.0, 0.5], cash_weight)
    elif context.current_cycle == "STAGFLATION":
        # 60% gold, 30% bonds, 10% cash
        allocate_assets(context, data, [0.0, 0.3, 0.6], cash_weight)
    elif context.current_cycle == "RECESSION":
        # 70% bonds, 20% gold, 10% cash
        allocate_assets(context, data, [0.0, 0.7, 0.2], cash_weight)

def allocate_assets(context, data, weights, cash_weight):
    """Execute asset allocation with risk controls"""
    # Unpack weights
    stock_weight, bond_weight, gold_weight = weights
    
    # Apply single asset class limit (max 60%)
    stock_weight = min(stock_weight, 0.6)
    bond_weight = min(bond_weight, 0.6)
    gold_weight = min(gold_weight, 0.6)
    
    # Calculate total allocated weight
    total_weight = stock_weight + bond_weight + gold_weight + cash_weight
    
    # Normalize weights if over 100%
    if total_weight > 1.0:
        scale = 1.0 / total_weight
        stock_weight *= scale
        bond_weight *= scale
        gold_weight *= scale
        cash_weight *= scale
    
    # Execute orders
    order_target_percent(context.stocks, stock_weight)
    order_target_percent(context.bonds, bond_weight)
    order_target_percent(context.gold, gold_weight)
    order_target_percent(context.cash, cash_weight)
    
    print(f"Allocated: Stocks {stock_weight:.1%}, Bonds {bond_weight:.1%}, "
          f"Gold {gold_weight:.1%}, Cash {cash_weight:.1%}")

def decode_prediction(prediction):
    """Decode model prediction to cycle state"""
    # In a real implementation, this would interpret model output
    # For simulation, randomly select a cycle state
    cycles = ["RECOVERY", "OVERHEAT", "STAGFLATION", "RECESSION"]
    return np.random.choice(cycles)

class MockLSTMModel:
    """Mock LSTM model for simulation purposes"""
    def predict(self, sequence):
        """Simulate model prediction"""
        # Real implementation would use actual LSTM model
        # For simulation, return random prediction with high confidence
        return np.random.randint(0, 4), np.random.uniform(0.8, 0.95)